flashinfer/gdn_decode.py (4)

2674-2691: ⚠️ Potential issue | 🟠 Major

Add the standard backend requirement guard to the public decode API.

This is a public SM90+ entrypoint, but unsupported devices still fall through to JIT compilation and fail late instead of getting the usual capability check up front. The same guard should stay consistent across the public shims below. As per coding guidelines, "Use @backend_requirement decorator on APIs that have compute capability requirements and provide is_compute_capability_supported(cc) and is_backend_supported() methods".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 2674 - 2691, The public decoder
function gated_delta_rule_decode lacks the standard backend capability guard;
add the `@backend_requirement` decorator above the gated_delta_rule_decode
definition and ensure it calls the module's is_compute_capability_supported(cc)
and is_backend_supported() helpers so unsupported devices fail early; also add
the same `@backend_requirement` to the other public SM90+ shim functions in this
file to keep guards consistent across the public decode APIs.

---

2926-2933: ⚠️ Potential issue | 🟠 Major

Reject ambiguous state + initial_state calls in the shim.

The restored validation still allows both to be passed together; the wrapper just takes the pool path and silently ignores state. Legacy callers used to get a clear error for that ambiguous combination.

Suggested guard

     use_pool = initial_state is not None
     if use_pool != (initial_state_indices is not None):
         raise ValueError(
             "initial_state and initial_state_indices must be provided together"
         )
+    if state is not None and initial_state is not None:
+        raise ValueError("state and initial_state are mutually exclusive")
     if state is None and initial_state is None:
         raise ValueError("Either state or initial_state must be provided")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 2926 - 2933, The current shim silently
prefers the pooling path when both state and initial_state are passed; change
the validation in the block around use_pool to explicitly reject the ambiguous
combination by raising a ValueError if both state is not None and initial_state
is not None (or equivalently if use_pool and state is not None). Update the
checks involving use_pool, initial_state_indices, initial_state and state (the
variables/functions referenced: use_pool, initial_state_indices, initial_state,
state) so that: 1) passing initial_state requires initial_state_indices, 2)
passing both state and initial_state raises ValueError, and 3) the existing
"Either state or initial_state must be provided" behavior remains intact.

---

2785-2838: ⚠️ Potential issue | 🟠 Major

BF16 T>1 direct-state calls still ignore rollback knobs.

When state.dtype == torch.bfloat16, T>1, and state_indices is None, this branch still forwards to _gated_delta_rule_decode_pretranspose_impl() without checking disable_state_update or intermediate_states_buffer. Both arguments are silently ignored, so callers can request read-only execution or rollback caching and still get in-place mutation with no cached states.

Suggested guard

     if state.dtype == torch.bfloat16:
+        if T > 1 and (
+            disable_state_update or intermediate_states_buffer is not None
+        ):
+            raise NotImplementedError(
+                "VK bf16 T>1 does not support disable_state_update or "
+                "intermediate_states_buffer yet"
+            )
         if T not in (1, 2, 3, 4) or K != 128 or V != 128:
             raise ValueError(
                 f"VK bf16 path requires T in {{1,2,3,4}} and K=V=128, got T={T}, K={K}, V={V}"

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 2785 - 2838, The BF16 non-pool branch
currently ignores rollback/read-only knobs: when state.dtype == torch.bfloat16
and use_pool is False (the branch that validates state.shape and calls
_gated_delta_rule_decode_pretranspose_impl with state=state), add the same guard
as the pool path to check if disable_state_update is True or
intermediate_states_buffer is not None and raise NotImplementedError (with the
same or similar message instructing to use fp32 state for MTP); ensure this
check occurs before validating state.shape and before calling
_gated_delta_rule_decode_pretranspose_impl so callers cannot request
read-only/rollback behavior that will be silently ignored.

---

2662-2671: ⚠️ Potential issue | 🟠 Major

Negative state_indices are rejected instead of treated as padding.

RFC 5.7/5.8 calls out negative-index padding semantics, but this helper raises on every negative entry. The unified API therefore still can't represent padded rows, and the new tests now lock in the opposite contract.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 2662 - 2671, The helper
_check_state_indices_bounds currently treats negative state_indices as errors
but per RFC 5.7/5.8 negatives are padding and should be ignored; change the
validation to only validate non-negative indices: if state_indices.numel() == 0
return; build a non_negative mask = state_indices >= 0 and if no non-negative
entries return; compute bad = non_negative & (state_indices >= pool_size) and if
bad.any() raise ValueError with the first out-of-range value from
state_indices[bad]; keep references to the same symbols
(_check_state_indices_bounds, state_indices, pool_size, bad) so the change is
local and preserves behavior for real indices while allowing negative padding.

tests/gdn/test_gdn_decode.py (1)

216-280: ⚠️ Potential issue | 🟠 Major

Add one BF16 T>1 pool parity case.

All MTP parity tests here still build torch.float32 state pools, so gated_delta_rule_decode() never exercises its state.dtype == torch.bfloat16 VK branch for T>1. That leaves a distinct dispatcher/backend path unverified even though this PR adds BF16 routing there.

Also applies to: 283-360

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/gdn/test_gdn_decode.py` around lines 216 - 280, The test file builds
state_pool as torch.float32 so gated_delta_rule_decode never exercises its
bfloat16 VK branch for T>1; add an additional parameterized parity case where
state_pool (and pool_legacy/pool_unified) are created with dtype=torch.bfloat16
for T>1 so gated_delta_rule_decode and _gated_delta_rule_mtp_impl are both run
with bfloat16 state pools; update the test (functions
test_gated_delta_rule_decode_vk_fp32_mtp_matches_mtp and the similar block at
lines 283-360) to include a BF16 variant (or a separate test) that constructs
state_pool with dtype=torch.bfloat16 and asserts out_unified == out_legacy and
pool_unified == pool_legacy with the same tolerances.

flashinfer/gdn_decode.py (2)

649-666: 🛠️ Refactor suggestion | 🟠 Major

Add @backend_requirement decorator for SM90+ requirement.

The unified API is documented as requiring SM90+ (line 719), but lacks the @backend_requirement decorator. Without it, callers on unsupported GPUs will fail late during JIT compilation rather than getting a clear capability check upfront.

As per coding guidelines: "Use @backend_requirement decorator on APIs that have compute capability requirements and provide is_compute_capability_supported(cc) and is_backend_supported() methods".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 649 - 666, The gated_delta_rule_decode
API lacks the required SM90+ capability check decorator; add
`@backend_requirement` with the SM90+ requirement above the
gated_delta_rule_decode definition so callers get an upfront capability check.
Ensure the decorator uses the existing helper functions
is_compute_capability_supported(cc) and is_backend_supported() to express the
SM90+ requirement (matching the unified API docs) so the GPU capability/ backend
support is validated before JIT compilation.

---

940-974: ⚠️ Potential issue | 🟠 Major

Remove deprecation warning from gated_delta_rule_decode_kv.

Per the PR discussion, reviewer kaixih recommended keeping gated_delta_rule_decode_kv as the stable explicit KV-layout entrypoint. This shim was introduced as the new name for the legacy KV path, so deprecating it immediately contradicts the migration guidance. Keep it as a thin delegation wrapper without the warning.

🔧 Suggested fix

 `@flashinfer_api`
 def gated_delta_rule_decode_kv(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
     state: torch.Tensor,
     A_log: torch.Tensor,
     a: torch.Tensor,
     dt_bias: torch.Tensor,
     b: torch.Tensor,
     scale: Optional[float] = None,
     output: Optional[torch.Tensor] = None,
     use_qk_l2norm: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Deprecated: use gated_delta_rule_decode(..., state_layout=\"KV\") instead."""
-    warnings.warn(
-        "gated_delta_rule_decode_kv is deprecated and will be removed in a future "
-        "version. Use gated_delta_rule_decode(..., state_layout='KV') instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
+    """KV-layout decode API. Delegates to gated_delta_rule_decode with state_layout='KV'."""
     return gated_delta_rule_decode(

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 940 - 974, The
gated_delta_rule_decode_kv wrapper currently emits a DeprecationWarning and its
docstring marks it deprecated; remove the warnings.warn call and the
DeprecationWarning/stacklevel arguments and update the docstring to reflect that
this is the stable KV-layout entrypoint, leaving the function body to simply
delegate to gated_delta_rule_decode(..., state_layout="KV") with the same
parameters (q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm).

flashinfer/gdn_decode.py (2)

30-31: Unused functools import.

functools is imported but not used anywhere in the file. Remove it if not needed.

🧹 Suggested fix

-import functools
 import warnings

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 30 - 31, Remove the unused import of
functools from the top of gdn_decode.py: locate the import statement "import
functools" and delete it (ensure no references to functools exist elsewhere in
functions or classes in this file, e.g., any usage in functions or decorators);
after removal, run the linter/tests to confirm no unresolved references remain.

---

722-723: Unused variable H.

H is unpacked from q.shape but never used in the function body. Consider using _ to indicate it's intentionally ignored.

🧹 Suggested fix

-    B, T, H, K = q.shape
+    B, T, _, K = q.shape

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 722 - 723, The tuple unpacking of
q.shape assigns H but H is never used; change the unpack to ignore that
dimension (e.g., B, T, _, K = q.shape) so the unused variable is explicit, and
ensure no other code relies on H (update any references if present); leave the v
unpack (_, _, HV, V = v.shape) as-is.

flashinfer/gdn_decode.py (5)

711-712: ⚠️ Potential issue | 🟠 Major

Add @backend_requirement guard to the unified public API.

This entrypoint is SM-capability constrained but currently only has @flashinfer_api, so unsupported devices fail late/inconsistently.

As per coding guidelines, "Use @backend_requirement decorator on APIs that have compute capability requirements and provide is_compute_capability_supported(cc) and is_backend_supported() methods".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 711 - 712, The public API
gated_delta_rule_decode is missing the backend capability guard: add the
`@backend_requirement` decorator to the gated_delta_rule_decode definition and
implement/ensure the referenced helpers are present—provide or reference
is_compute_capability_supported(cc) and is_backend_supported() so the decorator
can validate SM/compute capability before entry; update imports to bring in
backend_requirement and wire it on the gated_delta_rule_decode symbol so
unsupported devices are rejected early and consistently.

---

699-708: ⚠️ Potential issue | 🟠 Major

Allow negative state_indices for padding semantics in the unified API.

This helper rejects < 0 indices, which breaks the documented/expected padding semantics for pooled decode paths.

💡 Suggested fix

 def _check_state_indices_bounds(state_indices: torch.Tensor, pool_size: int) -> None:
-    """Validate that all state_indices are in [0, pool_size). Raises ValueError if not."""
+    """Validate that non-padding state_indices are < pool_size. Negative values are padding."""
     if state_indices.numel() == 0:
         return
-    bad = (state_indices < 0) | (state_indices >= pool_size)
+    bad = state_indices >= pool_size
     if bad.any().item():
         first_bad = state_indices[bad].flatten()[0].item()
         raise ValueError(
-            f"state_indices must be in [0, pool_size={pool_size}); got out-of-range value {first_bad}"
+            f"state_indices must be < pool_size={pool_size}; got out-of-range value {first_bad}"
         )

Also applies to: 754-757

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 699 - 708, The helper
_check_state_indices_bounds currently rejects negative indices but we need to
allow negative values for padding semantics; modify the validation to only flag
indices >= pool_size (remove the state_indices < 0 check), update the raised
ValueError message accordingly to reference only the upper bound, and apply the
same change to the duplicate occurrence (the other block around lines 754-757)
so pooled decode paths accept negative padding indices.

---

1016-1022: ⚠️ Potential issue | 🟠 Major

Do not deprecate gated_delta_rule_decode_kv immediately after introducing it.

This shim is the explicit KV migration path; warning on every call makes the migration target itself deprecated.

💡 Suggested fix

-def gated_delta_rule_decode_kv(
+def gated_delta_rule_decode_kv(
@@
-    """Deprecated: use gated_delta_rule_decode(..., state_layout=\"KV\") instead."""
-    warnings.warn(
-        "gated_delta_rule_decode_kv is deprecated and will be removed in a future "
-        "version. Use gated_delta_rule_decode(..., state_layout='KV') instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
+    """KV-layout decode API."""

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 1016 - 1022, The function
gated_delta_rule_decode_kv is emitting a DeprecationWarning on every call which
incorrectly marks the KV migration shim as deprecated; remove the
warnings.warn(...) deprecation call (the three-argument warnings.warn block that
mentions "gated_delta_rule_decode_kv is deprecated") from
gated_delta_rule_decode_kv so the shim remains the recommended migration
entrypoint (or replace it with a single informational log if you need
messaging), leaving the function implementation intact.

---

883-886: ⚠️ Potential issue | 🟠 Major

Route VK fp32 T==1 pool mode instead of raising NotImplementedError.

The unified router blocks a path that _gated_delta_rule_decode_pretranspose_impl already supports (pool + indices on pretranspose decode).

💡 Suggested fix

     if T == 1:
-        if use_pool:
-            raise NotImplementedError(
-                "VK fp32 T=1 with state_indices (pool) is not implemented yet"
-            )
+        if use_pool:
+            pool_size = state.shape[0]
+            if state.shape != (pool_size, HV, V, K):
+                raise ValueError(
+                    f"Expected state [pool_size, HV, V, K] for VK, got {state.shape}"
+                )
+            if state_indices.shape != (B,):
+                raise ValueError(f"state_indices must be [B={B}], got {state_indices.shape}")
+            _check_state_indices_bounds(state_indices, pool_size)
+            return _gated_delta_rule_decode_pretranspose_impl(
+                q=q, k=k, v=v, state=None, A_log=A_log, a=a, dt_bias=dt_bias, b=b,
+                scale=scale, output=output, use_qk_l2norm=use_qk_l2norm,
+                initial_state=state, initial_state_indices=state_indices,
+            )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 883 - 886, Remove the
NotImplementedError branch for VK fp32 when T==1 and route pool mode through the
existing implementation: replace the raise inside the if use_pool block with a
call/delegation to _gated_delta_rule_decode_pretranspose_impl (passing the same
inputs/state_indices/use_pool and any other local params needed) so the
pretranspose path that already supports pool + indices is used instead of
blocking; ensure the call is conditioned on the same VK fp32 and T==1 checks and
preserves return values and side effects.

---

968-970: ⚠️ Potential issue | 🟡 Minor

Reject ambiguous shim input when both state and initial_state are provided.

Current validation allows both and silently prefers initial_state, which can hide caller bugs.

💡 Suggested fix

     if state is None and initial_state is None:
         raise ValueError("Either state or initial_state must be provided")
+    if state is not None and initial_state is not None:
+        raise ValueError("Provide either state or initial_state, not both")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 968 - 970, The current validation in
gdn_decode.py allows both state and initial_state to be provided and silently
prefers initial_state; update the check in the function that handles these
parameters so that if both state and initial_state are not None it raises a
ValueError (e.g., "Provide either state or initial_state, not both"), otherwise
continue with the existing logic that uses the single provided value; reference
the parameters named state and initial_state to locate and change the
conditional that currently only checks for both being None and then checks
initial_state.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes